import os

from data.readData import TrainDataReader
from environment.electric_scheduling import PowerDayAheadSchedule
from environment.storage import Storage
import utils.show_policy as tool
# from model.DDPG.v1.trainer import Trainer
from model.SAC.v2.sac_trainer import sac_trainer as Trainer

trainData = TrainDataReader()
days = [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
MODEL_PATH = "./result/sac_model.pth"
fireStationNum = 10
trainDuration = 24 * 1

if __name__ == '__main__':
    # 负载加上外输电流即为需求 + trainData.Connection[2, 4] * trainData.ConnectionGraph[:24, 2]
    requirement = trainData.Payload[:trainDuration, 2]
    # # 实际的风光电数据
    windPower = trainData.NewPower[:trainDuration, 3]
    sunPower = trainData.NewPower[:trainDuration, 5]
    fireInfo = trainData.FireElectric[:fireStationNum, 4:9]
    storageInfo = trainData.StoragePower[:2, :]
    savingStorage = Storage(storageInfo[0])
    batteryStorage = Storage(storageInfo[1])
    fireLimit = []
    for i in range(fireStationNum):
        fireLimit.append([
            fireInfo[i][0] * fireInfo[i][1],
            fireInfo[i][0] * fireInfo[i][2],
            fireInfo[i][0] * fireInfo[i][3],
            fireInfo[i][0] * fireInfo[i][4],
        ])
    waterInfoP = trainData.WaterElectric[:50, 3]
    waterInfoG = trainData.WaterElectric[:50, 7]
    waterGraph = trainData.WaterPowerGraph
    waterLimit = []
    for i in range(len(waterInfoP)):
        graph_id = waterInfoG[i]
        for j in waterGraph:
            if j[0] == graph_id:
                station_limit = []
                for m in range(12):
                    for d in range(days[m]):
                        for h in range(24):
                            station_limit.append(waterInfoP[i] * j[m + 1])
                waterLimit.append(station_limit)
                break

    station_info = {
        "water-station-num": 50,
        "fire-station-num": fireStationNum
    }
    temp = [[waterLimit[i][j]
             for j in range(trainDuration)] for i in range(len(waterLimit))]
    # handler = Solver(requirement, fireLimit, [], fireInfo, sunPower, windPower)
    # handler.solve()
    env = PowerDayAheadSchedule(
        d=trainDuration,
        require_=requirement,
        fireP=fireLimit,
        fireExtra=trainData.FireElectric[:fireStationNum, 18],
        windP=windPower,
        sunP=sunPower,
        waterP=temp,
        saving=savingStorage,
        battery=batteryStorage,
        **station_info
    )
    if os.path.exists(MODEL_PATH):
        # agent = tool.LoadDDPGModel(env, MODEL_PATH)
        agent = tool.LoadSACV2Model(env, MODEL_PATH)
    else:
        agent, _ = Trainer(env)
    tool.ShowPolicy(env, agent)
